import numpy as np 
import cv2 
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import dicomLoad
# Step 1: Simulate 3D coronary artery data
def simulate_vessel():
    t = np.linspace(0, 4 * np.pi, 200)  # 增加点数以提高稳定性
    r = 5
    k = 0.5
    x = r * np.cos(t)
    y = r * np.sin(t)
    z = k * t + 10  # 确保 z 始终大于 0
    return np.vstack((x, y, z)).T
 
# Step 2: Project 3D points to 2D images
def project_points(points_3d, K, R, t):
    extrinsic = np.hstack((R, t))
    points_2d = K @ extrinsic @ np.vstack((points_3d.T, np.ones((1, points_3d.shape[0]))))
    points_2d = points_2d[:2] / points_2d[2]
    return points_2d.T
 
# Step 3: Add Gaussian noise 
def add_noise(points, sigma=0.1):  # 进一步减小噪声强度
    noise = np.random.normal(0, sigma, points.shape)
    return points + noise
 
# Step 4: Filter valid points
def filter_valid_points(points_2d_1, points_2d_2, img_shape):
    valid_points_1 = []
    valid_points_2 = []
    for pt1, pt2 in zip(points_2d_1, points_2d_2):
        if (np.isfinite(pt1[0]) and np.isfinite(pt1[1]) and
            np.isfinite(pt2[0]) and np.isfinite(pt2[1]) and
            0 <= pt1[0] < img_shape[1] and 0 <= pt1[1] < img_shape[0] and
            0 <= pt2[0] < img_shape[1] and 0 <= pt2[1] < img_shape[0]):
            valid_points_1.append(pt1)
            valid_points_2.append(pt2)
    
    # Convert to numpy arrays and ensure shape is (N, 2)
    points_2d_1 = np.array(valid_points_1).reshape(-1, 2)
    points_2d_2 = np.array(valid_points_2).reshape(-1, 2)
    
    # Check if there are any valid points
    if points_2d_1.size == 0 or points_2d_2.size == 0:
        raise ValueError("Error: No valid points remain after filtering.")
    
    return points_2d_1, points_2d_2 
 
# Step 5: Directly use simulated points for matching
def estimate_fundamental_matrix_direct(points_2d_1, points_2d_2):
    # Ensure the shape of points is (N, 2)
    points_2d_1 = points_2d_1.reshape(-1, 2)
    points_2d_2 = points_2d_2.reshape(-1, 2)
    
    # Ensure the data type is float32
    points_2d_1 = points_2d_1.astype(np.float32)
    points_2d_2 = points_2d_2.astype(np.float32)
    
    # Check if there are enough points 
    if len(points_2d_1) < 8 or len(points_2d_2) < 8:
        raise ValueError("Error: Not enough points to estimate fundamental matrix (minimum 8 points required).")
    
    # Use all points as matches 
    pts1 = points_2d_1
    pts2 = points_2d_2
    
    # Estimate fundamental matrix using RANSAC 
    F, mask = cv2.findFundamentalMat(pts1, pts2, cv2.FM_RANSAC)
    
    # # Filter inliers 
    # if mask is not None:
    #     pts1 = pts1[mask.ravel() == 1]
    #     pts2 = pts2[mask.ravel() == 1]
    # else:
    #     raise ValueError("Error: Fundamental matrix estimation failed. Check the input points.")
    
    return F, pts1, pts2
 
# Step 6: Triangulation 
def triangulate_points(pts1, pts2, K, R1, t1, R2, t2):
    P1 = K @ np.hstack((R1, t1))
    P2 = K @ np.hstack((R2, t2))
    points_4d_hom = cv2.triangulatePoints(P1, P2, pts1.T, pts2.T)
    points_3d = points_4d_hom[:3] / points_4d_hom[3]
    return points_3d.T 
#%%
# Main function 
if __name__ == "__main__":
    pass
    #%%
    import qca

    sigmoid_inputs = np.arange(0,250,0.01)

    sigmoid_outputs = []
    for p in sigmoid_inputs:
        sigmoid_outputs.append(qca.speed_function(p))
    print("Sigmoid Function Input :: {}".format(sigmoid_inputs))
    print("Sigmoid Function Output :: {}".format(sigmoid_outputs))

    plt.plot(sigmoid_inputs,sigmoid_outputs)
    plt.xlabel("Sigmoid Inputs")
    plt.ylabel("Sigmoid Outputs")
    plt.show()
    #%%
    points_2d_1 = np.loadtxt('../resource/key_points0.txt', usecols=(0, 1))
    points_2d_2 = np.loadtxt('../resource/key_points4.txt', usecols=(0, 1))
    
    dicom_path1 = "../resource/IM000000"  # 替换为你的DICOM文件路径
    params1 = dicomLoad.load_dicom_parameters(dicom_path1)
    plt.imshow(params1['image'][22,:,:],cmap='gray')
    gray_img=params1['image'][22,:,:]
    print(params1)
    # 计算内参矩阵
    K1 = dicomLoad.compute_intrinsic_matrix(params1)
    print("内参矩阵 K:\n", K1)
    # 计算外参矩阵
    R1,T1 = dicomLoad.compute_extrinsic_matrix(params1)
    points_2d_1=points_2d_1/params1['pixel_spacing'][0]
    points_2d_1= points_2d_1.astype(np.int16)
    points_2d_1[:,[0,1]]=points_2d_1[:,[1,0]]
    path1=[]
    for i in range(1,len(points_2d_1)):
        print(i)
        arrival_times = qca.wavefront_propagation(params1['image'][22,:,:], points_2d_1[i-1],points_2d_1[i])
        path0= qca.backtracking(arrival_times, points_2d_1[i])
        print(points_2d_1[i-1],points_2d_1[i])
        print(path0[0],path0[-1])
        path1+=path0
    plt.figure(figsize=(20,20))
    plt.scatter([p[1] for p in path1],[p[0] for p in path1],s=1,c=[0,1,0])
    plt.imshow(gray_img)
    #%%

#%%

    dicom_path2 = "../resource/IM000004"  # 替换为你的DICOM文件路径
    params2 = dicomLoad.load_dicom_parameters(dicom_path2)
    plt.imshow(params2['image'][20,:,:],cmap='gray')
    print(params2)
    # 计算内参矩阵
    K2 = dicomLoad.compute_intrinsic_matrix(params2)
    print("内参矩阵 K:\n", K2)
    # 计算外参矩阵
    R2,T2 = dicomLoad.compute_extrinsic_matrix(params2)
    gray_img=params2['image'][20,:,:]

    points_2d_2=points_2d_2/0.209062
    points_2d_2= points_2d_2.astype(np.int16)
    points_2d_2[:,[0,1]]=points_2d_2[:,[1,0]]
    path2=[]
    for i in range(1,len(points_2d_2)):
        print(i)
        arrival_times = qca.wavefront_propagation(gray_img, points_2d_2[i-1],points_2d_2[i])
        path0= qca.backtracking(arrival_times, points_2d_2[i])
        print(points_2d_2[i-1],points_2d_2[i])
        print(path0[0],path0[-1])
        path2+=path0
    plt.figure(figsize=(20,20))
    plt.scatter([p[1] for p in path2],[p[0] for p in path2],s=1,c=[0,1,0])
    plt.imshow(gray_img)
    #%%
    # Estimate fundamental matrix directly
    try:
        F, pts1, pts2 = estimate_fundamental_matrix_direct(np.array(path1), np.array(path2[0:386]))
    except ValueError as e:
        print(e)
        exit()
    
    # Triangulate points
    points_reconstructed = triangulate_points(pts1, pts2, K1, R1, T1, R2, T2)
    print(points_reconstructed)
    # Plot 3D reconstruction 
    fig = plt.figure()

    ax2 = fig.add_subplot(111, projection='3d')

    ax2.scatter(points_reconstructed[:, 0], points_reconstructed[:, 1], points_reconstructed[:, 2], c='g', marker='o')
    ax2.plot(points_reconstructed[:, 0], points_reconstructed[:, 1], points_reconstructed[:, 2], c='r', marker='o')

    plt.show()
# %%
